#ifndef YOLOV8_TENSORRT_H
#define YOLOV8_TENSORRT_H

#include <NvInfer.h>
#include <cuda_runtime.h>
#include <opencv2/opencv.hpp>
#include <vector>
#include <memory>
#include <string>

struct Detection {
    float x1, y1, x2, y2;
    float confidence;
    int classId;
};

class YOLOv8TensorRT {
public:
    YOLOv8TensorRT(const std::string& enginePath);
    ~YOLOv8TensorRT();
    
    std::vector<Detection> detect(const cv::Mat& image);
    
private:
    // TensorRT相关
    std::unique_ptr<nvinfer1::ICudaEngine> engine;
    std::unique_ptr<nvinfer1::IExecutionContext> context;
    cudaStream_t stream;
    
    // 缓冲区
    void* buffers[2];  // 输入和输出
    float* inputBuffer;
    float* outputBuffer;
    
    // 模型参数 - 需要根据您的模型修改
    const int INPUT_W = 640;      // ⚠️ 需要修改：输入宽度
    const int INPUT_H = 640;      // ⚠️ 需要修改：输入高度
    const int NUM_CLASSES = 80;   // ⚠️ 需要修改：类别数量（COCO=80）
    const int OUTPUT_SIZE = 8400; // ⚠️ 需要修改：输出大小 (对于640x640通常是8400)
    
    // 阈值参数 - 可根据需求调整
    float CONF_THRESHOLD = 0.85f;  // ⚠️ 可调整：置信度阈值
    float NMS_THRESHOLD = 0.45f;   // ⚠️ 可调整：NMS阈值
    
    // 辅助函数
    void preprocess(const cv::Mat& image, float* inputBuffer);
    std::vector<Detection> postprocess(float* output, int imgWidth, int imgHeight);
    void nms(std::vector<Detection>& detections);
    bool loadEngine(const std::string& enginePath);
};

// Logger for TensorRT
class Logger : public nvinfer1::ILogger {
    void log(Severity severity, const char* msg) noexcept override {
        if (severity <= Severity::kWARNING)
            std::cout << msg << std::endl;
    }
};

#endif